package testing

import (
	"context"
	"fmt"
	"time"

	"github.com/louloulin/dataflare/pkg/workflow"
	"github.com/louloulin/dataflare/pkg/workflow/connector"
)

// TestFramework represents a workflow testing framework
type TestFramework struct {
	// engine is the workflow engine
	engine *workflow.Engine
	// mockSources is a map of mock source connectors
	mockSources map[string]*MockSourceConnector
	// mockSinks is a map of mock sink connectors
	mockSinks map[string]*MockSinkConnector
	// mockProcessors is a map of mock processors
	mockProcessors map[string]*MockProcessor
}

// NewTestFramework creates a new workflow testing framework
func NewTestFramework() (*TestFramework, error) {
	// Create a workflow engine
	engine, err := workflow.NewEngine()
	if err != nil {
		return nil, fmt.Errorf("failed to create workflow engine: %w", err)
	}

	return &TestFramework{
		engine:         engine,
		mockSources:    make(map[string]*MockSourceConnector),
		mockSinks:      make(map[string]*MockSinkConnector),
		mockProcessors: make(map[string]*MockProcessor),
	}, nil
}

// RegisterMockSource registers a mock source connector
func (f *TestFramework) RegisterMockSource(name string, data []interface{}) {
	f.mockSources[name] = NewMockSourceConnector(data)
}

// RegisterMockSink registers a mock sink connector
func (f *TestFramework) RegisterMockSink(name string) *MockSinkConnector {
	sink := NewMockSinkConnector()
	f.mockSinks[name] = sink
	return sink
}

// RegisterMockProcessor registers a mock processor
func (f *TestFramework) RegisterMockProcessor(name string, processor func(record *connector.Record) *connector.Record) {
	f.mockProcessors[name] = NewMockProcessor(processor)
}

// DeployWorkflow deploys a workflow
func (f *TestFramework) DeployWorkflow(definition *workflow.Definition) (string, error) {
	// Replace sources with mock sources
	for name, source := range definition.Sources {
		if mockSource, ok := f.mockSources[name]; ok {
			source.Type = "mock"
			source.Config = map[string]interface{}{
				"mock_id": mockSource.ID,
			}
		}
	}

	// Replace sinks with mock sinks
	for name, sink := range definition.Sinks {
		if mockSink, ok := f.mockSinks[name]; ok {
			sink.Type = "mock"
			sink.Config = map[string]interface{}{
				"mock_id": mockSink.ID,
			}
		}
	}

	// Replace processors with mock processors
	for name, processor := range definition.Processors {
		if mockProcessor, ok := f.mockProcessors[name]; ok {
			processor.Type = "mock"
			processor.Config = map[string]interface{}{
				"mock_id": mockProcessor.ID,
			}
		}
	}

	// Deploy the workflow
	return f.engine.DeployWorkflow(context.Background(), definition)
}

// RunWorkflow runs a workflow
func (f *TestFramework) RunWorkflow(workflowID string, parameters map[string]interface{}) (string, error) {
	// Start the workflow
	runID, err := f.engine.StartWorkflow(context.Background(), workflowID, parameters)
	if err != nil {
		return "", fmt.Errorf("failed to start workflow: %w", err)
	}

	// Wait for the workflow to complete
	for {
		// Get workflow status
		status, err := f.engine.GetWorkflowStatus(context.Background(), runID)
		if err != nil {
			return "", fmt.Errorf("failed to get workflow status: %w", err)
		}

		// Check if the workflow is completed
		if status.Status == "completed" || status.Status == "failed" {
			return runID, nil
		}

		// Wait for a short time
		time.Sleep(100 * time.Millisecond)
	}
}

// GetSinkRecords gets records from a mock sink
func (f *TestFramework) GetSinkRecords(name string) ([]*connector.Record, error) {
	// Get the mock sink
	sink, ok := f.mockSinks[name]
	if !ok {
		return nil, fmt.Errorf("mock sink %s not found", name)
	}

	// Get records from the sink
	return sink.GetRecords(), nil
}

// Shutdown shuts down the framework
func (f *TestFramework) Shutdown() error {
	// Shutdown the engine
	return f.engine.Shutdown()
}
